# -*- coding: UTF-8 -*-
import sys

import dateutil
import pytz
from pymongo import MongoClient
from datetime import datetime, timezone
import pandas as pd
import time
import traceback
import numpy as np
from core.xgboost_model import XGBoostModel
import requests
import json
from pandas.plotting import scatter_matrix
from matplotlib import pyplot as plt

desired_width = 320
pd.set_option('display.width', desired_width)
pd.set_option('display.max_columns', 20)
pd.set_option('display.max_rows', 200)


def gen_signal(window=60, seq_len=72, predict_signal_name='return'):
    print('---------------------------------- begin ----------------------------------')
    print('predict_type: %s, window: %d, seq_len: %d ' % (predict_signal_name, window, seq_len))
    predict_db = MongoClient('mongodb://39.104.227.148:27019/')['predict_signal']
    predict_iter = predict_db['%s_%s_%s' % (predict_signal_name, window, seq_len)].find().sort('_id', 1)
    predict_list = list(predict_iter)
    df = pd.DataFrame(predict_list)

    df['time'] = pd.to_datetime(df['_id'], unit='s')
    df = df.get(['_id', 'time', 'predict_result', 'predict_%s' % predict_signal_name])

    begin_timestamp = int(df.iloc[0]['_id'])
    end_timestamp = int(df.iloc[-1]['_id'])

    price_param = {'_id': {'$gte': begin_timestamp * 1000, '$lte': end_timestamp * 1000}}
    print('price_param: %s' % price_param)
    db = MongoClient('mongodb://39.104.227.148:27018/')['exchange']
    price_iter = db['okex_btc_usdt'].find(price_param).sort('_id', 1)
    price_list = list(price_iter)
    price_df = pd.DataFrame(price_list)
    price_df['time'] = pd.to_datetime(price_df['_id'], unit='ms')

    if predict_signal_name == 'return':
        price_df = price_df[::window]
        price_df.reset_index(inplace=True, drop=True)
        price_df['real_return'] = (price_df['close'].shift(-1) - price_df['close']) / price_df['close']
        price_df = price_df.get(['_id', 'time', 'real_return'])
    elif predict_signal_name == 'volatility':
        price_df['real_%s' % predict_signal_name] = price_df['close'].rolling(window).std().shift(-window)

    result = pd.merge(df, price_df, left_on='time', right_on='time', suffixes=('_predict', '_price'))
    # result = result.get(['time', 'real_return', 'predict_return'])
    result = result.get(['time', 'real_%s' % predict_signal_name, 'predict_%s' % predict_signal_name])
    print(result.corr())
    result_analysis(result, 'real_%s' % predict_signal_name, 'predict_%s' % predict_signal_name)
    result.to_csv('predict_%s_1201_1219_%s_v2.csv' % (predict_signal_name, window), index=False)


def result_analysis(df, real_y_column, predict_y_column):
    threshold = sorted(df[real_y_column].abs(), reverse=True)[int(len(df) * 0.2)]
    threshold_df = df.loc[abs(df[real_y_column]) >= threshold]
    print('\nthreshold is %s, percentage is %s' % (threshold, (float(len(threshold_df)) / len(df))))

    predict_avg_all = abs(df[predict_y_column]).mean()
    predict_threshold_avg = abs(threshold_df[predict_y_column]).mean()
    print('avg_all: %s' % round(predict_avg_all, 4))
    print('avg: %s' % round(predict_threshold_avg, 4))
    print('result: %s' % round((predict_threshold_avg / predict_avg_all), 4))

    real_avg_all = abs(df[real_y_column]).mean()
    real_threshold_avg = abs(threshold_df[real_y_column]).mean()
    print('avg_all: %s' % round(real_avg_all, 4))
    print('avg: %s' % round(real_threshold_avg, 4))
    print('result: %s' % round((real_threshold_avg / real_avg_all), 4))


def export_hist_signal():
    # coll_list = ['test_return_15_72', 'test_return_30_72', 'test_return_60_72', 'test_return_15_132',
    #              'test_return_30_132', 'test_volatility_15_72', 'test_volatility_30_72']
    coll_list = ['return_15_72', 'return_30_72', 'return_60_72', 'return_15_132',
                 'return_30_132', 'volatility_15_72', 'volatility_30_72']
    # coll_list = ['volatility_15_72', 'volatility_30_72']
    predict_db = MongoClient('mongodb://39.104.227.148:27019/')['predict_signal']
    for c in coll_list:
        it = predict_db[c].find().sort('_id', 1)
        l = list(it)
        df = pd.DataFrame(l)
        # df.rename_axis()
        name = './saved_predict_result/export_signal/' + 'export_' + c.replace('test', 'export') + '.csv'

        predict_column_name = None
        if 'return' in c:
            predict_column_name = 'predict_return'
        if 'volatility' in c:
            predict_column_name = 'predict_volatility'
        df = df.get(['_id', predict_column_name])
        df = df.round(5)
        df['utc_time'] = pd.to_datetime(df['_id'], unit='s')

        # df = df.get(['utc_time', '_id', predict_column_name])
        df = df.get(['_id', predict_column_name])
        df.rename(columns={'_id': 'timestamp'}, inplace=True)
        df = df[(df['timestamp'] >= 1545300000)]
        df.to_csv(name, index=False)


def validate_price():
    # begin_utc_time = '2018-12-25 00:00:00'
    # end_utc_time = '2018-12-24 11:30:00'

    begin_local_time = '2018-12-25 14:30:00'
    end_local_time = '2018-12-25 15:30:00'
    # begin_timestamp = int(datetime.strptime(begin_local_time, '%Y-%m-%d %H:%M:%S').timestamp())
    # end_timestamp = int(datetime.strptime(end_local_time, '%Y-%m-%d %H:%M:%S').timestamp())
    begin_timestamp = 1545465600
    end_timestamp = 1545721200

    # begin_timestamp = int(
    #     datetime.strptime(begin_utc_time, '%Y-%m-%d %H:%M:%S').replace(tzinfo=timezone.utc).timestamp())
    # end_timestamp = int(
    #     datetime.strptime(end_utc_time, '%Y-%m-%d %H:%M:%S').replace(tzinfo=timezone.utc).timestamp())
    #
    price_param = {'_id': {'$gte': begin_timestamp * 1000, '$lte': end_timestamp * 1000}}
    print('price_param: %s' % price_param)
    db = MongoClient('mongodb://39.104.227.148:27018/')['exchange']
    price_iter = db['okex_btc_usdt'].find(price_param).sort('_id', 1)
    price_list = list(price_iter)
    print(price_list)
    df = pd.DataFrame(price_list)
    df['time'] = pd.to_datetime(df['_id'], unit='ms')
    print(df)


# def converter(row, local_time_column_name='local_time', utc_time_column_name='time'):
#     return row['time'].tz_localize('UTC').tz_convert(row['UTC'])


def compare_with_stock_pearson():
    utc_begin_time = '2018-12-01 00:00:00'
    utc_end_time = '2018-12-26 6:00:00'
    df = pd.read_csv('/Users/liujingjun/work/ml/LSTM/data/stock/IF000_hour.csv')
    df['utc_time'] = pd.to_datetime(df['time']).dt.tz_localize('Asia/Shanghai').dt.tz_convert('UTC')
    df = df[(df['utc_time'] >= utc_begin_time) & (df['utc_time'] <= utc_end_time)]
    df['timestamp'] = df['utc_time'].values.astype(np.int64) // 10 ** 9
    df = df.get(['timestamp', 'close'])
    df['return'] = (df['close'].shift(-1) - df['close']) / df['close']
    df.reset_index(inplace=True, drop=True)
    print(df)

    predict_db = MongoClient('mongodb://39.104.227.148:27019/')['predict_signal']
    return_table = predict_db['test_return_60_72']

    begin_timestamp = int(
        datetime.strptime(utc_begin_time, '%Y-%m-%d %H:%M:%S').replace(tzinfo=timezone.utc).timestamp())
    end_timestamp = int(
        datetime.strptime(utc_end_time, '%Y-%m-%d %H:%M:%S').replace(tzinfo=timezone.utc).timestamp())
    price_param = {'_id': {'$gte': begin_timestamp, '$lte': end_timestamp}}
    it = return_table.find(price_param).sort('_id', 1)
    return_df = pd.DataFrame(list(it))

    return_df = return_df.get(['_id', 'predict_return', 'next_price'])
    print(return_df)

    total_df = pd.merge(df, return_df, left_on='timestamp', right_on='_id')
    print(total_df[0:5])
    # total_df['next_price_shift'] = total_df['next_price'].shift(1)
    print(total_df[0:5])

    del total_df['timestamp']
    del total_df['_id']
    total_df.fillna(0, inplace=True)
    # total_df = total_df.round(3)

    # print(total_df)
    print(total_df.corr())
    total_df.get(['close', 'next_price']).plot(figsize=(20, 20))
    plt.show()

    df_norm = (total_df - total_df.min()) / (total_df.max() - total_df.min())
    df_norm['next_price_shift'] = df_norm['next_price'].shift(1)
    df_norm = df_norm.get(['close', 'next_price', 'next_price_shift'])
    df_norm.plot(figsize=(20, 20))
    # scatter_matrix(total_df, alpha=0.2, figsize=(10, 10), diagonal='kde')
    plt.show()
    pass


def export_online_signals():
    predict_type = 'return'
    window = 30
    sequence_length = 132
    table_name = '%s_%s_%s' % (predict_type, window, sequence_length)

    predict_db = MongoClient('mongodb://39.104.227.148:27019/')['predict_signal']
    # it = predict_db['return_15_132'].find()
    it = predict_db[table_name].find({'_id': {'$gte': 1545753600}})
    doc_list = list(it)
    df = pd.DataFrame(doc_list)
    if predict_type == 'return':
        df = df.get(['_id', 'predict_return'])
    elif predict_type == 'volatility':
        df = df.get(['_id', 'predict_volatility'])
        pass
    # df = df.get(['_id', 'predict_time', 'predict_return'])
    print(df)
    df = df.round(5)
    df.to_csv('predict_%s_20181226_20190103_30_v2.csv' % predict_type, index=False)


def compare_two_model_predict_result(predict_type='volatility', window=15, sequence_length=72):
    print('predict_type: %s, window: %d, sequence: %d' % (predict_type, window, sequence_length))
    table_name = '%s_%s_%s' % (predict_type, window, sequence_length)
    test_table_name = 'test_test_%s_%s_%s' % (predict_type, window, sequence_length)

    predict_db = MongoClient('mongodb://39.104.227.148:27019/')['predict_signal']
    test_it = predict_db[test_table_name].find().sort('_id', 1)
    test_row_list = list(test_it)
    test_df = pd.DataFrame(test_row_list)
    begin_time = int(test_df.iloc[0]['_id'])
    end_time = int(test_df.iloc[-1]['_id'])

    online_it = predict_db[table_name].find({'_id': {'$gte': begin_time, '$lte': end_time}}).sort('id', 1)
    online_row_list = list(online_it)
    online_df = pd.DataFrame(online_row_list)
    if predict_type == 'return':
        test_df = test_df.get(['_id', 'predict_return'])
        online_df = online_df.get(['_id', 'predict_return'])
    elif predict_type == 'volatility':
        test_df = test_df.get(['_id', 'predict_volatility'])
        online_df = online_df.get(['_id', 'predict_volatility'])

    price_param = {'_id': {'$gte': begin_time * 1000, '$lte': end_time * 1000}}
    db = MongoClient('mongodb://39.104.227.148:27018/')['exchange']
    price_iter = db['okex_btc_usdt'].find(price_param).sort('_id', 1)
    price_df = pd.DataFrame(list(price_iter))
    price_df['_timestamp'] = (price_df['_id'] / 1000).astype(int)
    price_df['time'] = pd.to_datetime(price_df['_id'], unit='ms')

    if predict_type == 'return':
        price_df = price_df[::window]
        price_df.reset_index(inplace=True, drop=True)
        price_df['real_return'] = (price_df['close'].shift(-1) - price_df['close']) / price_df['close']
        price_df = price_df.get(['_timestamp', 'real_return'])
    elif predict_type == 'volatility':
        price_df['real_%s' % predict_type] = price_df['close'].rolling(window).std().shift(-window)
        price_df = price_df.get(['_timestamp', 'real_volatility'])

    price_df.fillna(0, inplace=True)

    df = pd.merge(left=test_df, right=online_df, left_on='_id', right_on='_id', how='left',
                  suffixes=('_test', '_online'))
    df = pd.merge(left=df, right=price_df, left_on='_id', right_on='_timestamp', how='left')
    # print(df)
    df = df.get(['predict_%s_test' % predict_type, 'predict_%s_online' % predict_type, 'real_%s' % predict_type])
    print(df.corr())


if __name__ == '__main__':
    # gen_signal(window=15, seq_len=72, predict_signal_name='return')
    # gen_signal(window=30, seq_len=72, predict_signal_name='return')
    # gen_signal(window=60, seq_len=72, predict_signal_name='return')
    # gen_signal(window=15, seq_len=132, predict_signal_name='return')
    # gen_signal(window=30, seq_len=132, predict_signal_name='return')
    # gen_signal(window=15, seq_len=72, predict_signal_name='volatility')
    # gen_signal(window=30, seq_len=72, predict_signal_name='volatility')

    # result_analysis()
    export_hist_signal()
    # compare_with_stock_pearson()

    # d = datetime.strptime('2018-12-26 00:00:00', '%Y-%m-%d %H:%M:%S')
    # print(d.timestamp())
    # exit()

    # export_online_signals()
    # test()

    # compare_two_model_predict_result(predict_type='return', window=15, sequence_length=72)
    # compare_two_model_predict_result(predict_type='return', window=30, sequence_length=72)
    # compare_two_model_predict_result(predict_type='return', window=60, sequence_length=72)
    # compare_two_model_predict_result(predict_type='return', window=15, sequence_length=132)
    # compare_two_model_predict_result(predict_type='return', window=30, sequence_length=132)
    # compare_two_model_predict_result(predict_type='volatility', window=15, sequence_length=72)
    # compare_two_model_predict_result(predict_type='volatility', window=30, sequence_length=72)
    pass
